import torch
from eval import show_without_plt
import matplotlib.pyplot as plt
import numpy as np
# 中文显示
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['SimHei']
matplotlib.rcParams['axes.unicode_minus'] = False

def show_plot(load_tensors, is_rotated):
    if is_rotated:
        rota_time = "旋转后"
    else:
        rota_time = "旋转前"

    counts = []
    bins = []
    # 遍历加载的张量
    for tensor in load_tensors:
        temp = show_without_plt(tensor)
        counts.append(temp["counts"])
        bins = temp["bins"]
    # 计算平均值
    count = np.array(counts).mean(axis=0)
    # 转换区间为字符串形式
    bin_labels = [f"{bins[i]}" for i in range(len(bins))]
    print(bin_labels)

    # 绘制柱状图
    plt.figure(figsize=(12, 6))
    bars = plt.bar(bin_labels[:-1], count, color='skyblue', edgecolor='black',align= 'edge',width=1.0)

    # 手动设置 x 轴刻度以包括最后一个边界
    plt.xticks(np.arange(len(bin_labels)), bin_labels)

    plt.title(f"{rota_time}各block平均的featuremap数值分布", fontsize=15)

    # 设置横轴和纵轴标签
    plt.xlabel("区间", fontsize=12)
    plt.ylabel("统计数", fontsize=12)

    # 添加网格
    plt.grid(axis='y', linestyle='--', linewidth=0.5)

    # 在柱状图上标注数值
    for bar in bars:
        height = bar.get_height()  # 获取柱条高度
        plt.text(bar.get_x() + bar.get_width() / 2, height,  # 在柱条顶部居中显示
                 f'{height:.2f}',  # 格式化为两位小数
                 ha='center', va='bottom', fontsize=12, color='black')  # 水平居中，垂直靠下

    # 显示图形并保存
    plt.tight_layout()
    plt.show()

def show_mut_plot(load_tensors, is_rotated):
    counts = []
    bins = []

    if is_rotated:
        rota_time = "旋转后"
    else:
        rota_time = "旋转前"

    # 遍历加载的张量
    for tensor in load_tensors:
        temp = show_without_plt(tensor)
        counts.append(temp["counts"])
        bins = temp["bins"]


    # 转换区间为字符串形式
    bin_labels = [f"{bins[i]}" for i in range(len(bins))]
    print(bin_labels)

    fig, axs = plt.subplots(4, 8, figsize=(80, 20))  # 4行8列的子图

    for i, ax in enumerate(axs.flat):

        # 在特定子图上绘制柱状图
        bars = ax.bar(bin_labels[:-1], counts[i], color='skyblue', edgecolor='black', align='edge', width=1.0)

        ax.set_title(f"{rota_time}第{i+1}个block的featuremap数值分布", fontsize=15)



        # 设置横轴和纵轴标签
        ax.set_xlabel(f"区间", fontsize=12)
        ax.set_ylabel("统计数", fontsize=12)

        # 手动设置 x 轴刻度以包括最后一个边界
        # 手动设置 x 轴刻度和标签
        ax.set_xticks(np.arange(len(bin_labels)), bin_labels)


        # 添加网格
        ax.grid(axis='y', linestyle='--', linewidth=0.5)

        # 在柱状图上标注数值
        for bar in bars:
            height = bar.get_height()
            ax.text(bar.get_x() + bar.get_width() / 2, height,
                    f'{height:.0f}', ha='center', va='bottom', fontsize=10, color='black')
    plt.title("各block平均的featuremap数值分布", fontsize=15)
    plt.tight_layout()
    plt.show()

if __name__ == '__main__':
    # 加载数据
    load_tensors = torch.stack([torch.load(f"../ffn/second_line_input/tensor{i}.pth",weights_only=True) for i in range(32)])
    # show_plot(load_tensors, False)
    # show_mut_plot(load_tensors, False)

    # show_without_plt(load_tensors, print_info=True)
    # 设备
    rotate_matrix = torch.load("../1-1-2.pth", weights_only=True).to("cpu")


    load_tensors = torch.matmul(load_tensors,rotate_matrix)

    # show_without_plt(load_tensors, print_info=True)

    show_plot(load_tensors, True)
    show_mut_plot(load_tensors, True)